import cv2
import numpy as np
import os
from ultralytics import YOLO

# 文件路径
#video_path = '/home/wangziqi/runs/detect/predata/lic1.mp4'
video_path = input("请输入待检测视频的路径: ")
output_path = '/home/wangziqi/runs/detect/predict/blurred_two.mp4'
labels_path = '/home/wangziqi/runs/detect/predict/labels/'

# 初始化 YOLO 模型

model_path = '/home/wangziqi/double/labels/runs/detect/train/weights/best.pt'
model = YOLO(model_path)

# 加载视频
cap = cv2.VideoCapture(video_path)
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter(output_path, fourcc, int(cap.get(cv2.CAP_PROP_FPS)), 
                      (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))))

frame_num = 0

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break

    # YOLO 推断
    results = model.predict(frame)
    
    # 处理检测结果
    

    for box in results[0].boxes:
        if box.cls.item() in [ 0.0,1.0]:
            m = box.xyxy.cpu().tolist()

            y1, y2, x1, x2 = map(int, [m[0][1], m[0][3], m[0][0], m[0][2]])
            face_region = frame[y1:y2, x1:x2]
            blurred_face = cv2.GaussianBlur(face_region, (99, 99), 30)
            frame[y1:y2, x1:x2] = blurred_face


    print("处理后视频保存为：")
    print(output_path)

    out.write(frame)
    frame_num += 1

cap.release()
out.release()
